#!/usr/bin/python3
import sys
import os
import time
import logging

import numpy as np
from IPython import embed
from PySide2.QtCore import Qt, Slot, Signal
from PySide2.QtGui import QPainter, QIcon, QPixmap, QGuiApplication
from PySide2.QtWidgets import (QAction, QApplication, QFileDialog, QHBoxLayout,
                               QHeaderView, QLabel, QLineEdit, QMainWindow,
                               QPushButton, QSplitter, QTableWidget, QDesktopWidget,
                               QTableWidgetItem, QVBoxLayout, QWidget, QComboBox, QSplashScreen)
from vispy import scene
from vispy.scene import visuals
from vispy.scene.canvas import SceneCanvas
from vispy.geometry import generation
from vispy.io import read_mesh, write_mesh

from .MeshLoader import LoadMesh, LoadUrdf, PlotSphere, Path
from .RobotController import Robot


class MainWindow(QMainWindow):
    """
            RobotZen
    |----------------------------|
    |file|tools|                 |
    |----------------------------|
    | s |                        |
    | i |                        |
    | d |                        |
    | e |      canvas            |
    | b |                        |
    | a |                        |
    | r |                        |
    ------------------------------
    | status                     |
    ------------------------------

    TODO:
    Only GUI code should be kept here.

    Features:
    - add mesh visual node from obj/stl mesh file
    - scene graph tree
    - robot simulation support

    Next:
    - Fix robot FK/IK problem

    Plan:
    - select a visual node
    - pop up visual node property after selected
    - delete selected visual node
    - insert simple geometry visual node
    """
    LOGGER_COLORS = {logging.DEBUG: "grey", logging.INFO: "green",
                     logging.ERROR: "red", logging.WARNING: "yellow"}
    status_signal = Signal(str, int, name="status message")
    robot_target = Signal(np.ndarray, name="robot target joint positions")

    def __init__(self):
        super(MainWindow, self).__init__()
        self.setWindowTitle('RobotZen')
        self.setWindowIcon(QIcon('resources/robot.png'))
        size = QGuiApplication.primaryScreen().availableSize()
        self.setGeometry(0, 0, size.width(), size.height())

        # data
        self.robot = None
        self.canvas = Canvas()

        # Menu
        self.menu = self.menuBar()
        file_menu = self.menu.addMenu("File")

        # Open file
        open_action = QAction("Open", self)
        open_action.setShortcut("Ctrl+O")
        open_action.triggered.connect(self.open_file)
        file_menu.addAction(open_action)

        # Exit QAction
        exit_action = QAction("Exit", self)
        exit_action.setShortcut("Ctrl+Q")
        exit_action.triggered.connect(self.exit_app)
        file_menu.addAction(exit_action)

        geom_menu = self.menu.addMenu("Geometries")
        sphere = QAction("Sphere", self)
        sphere.triggered.connect(lambda: PlotSphere(self.canvas, 2.0))
        sphere.triggered.connect(lambda: self.sidebar.update_tree())
        geom_menu.addAction(sphere)

        debug_menu = self.menu.addMenu("Debug")
        debug = QAction("Embed", self)
        debug.triggered.connect(embed)
        debug_menu.addAction(debug)

        # central zone
        splitter = QSplitter(Qt.Horizontal)

        # side bar
        self.sidebar = SideBar(
            self.canvas.view, self.status_signal, self)
        splitter.addWidget(self.sidebar)

        self.canvas.create_native()
        self.canvas.native.setParent(self)
        splitter.addWidget(self.canvas.native)
        self.setCentralWidget(splitter)

        self.status = self.statusBar()
        self.status_lbl = QLabel()
        # self.status.addPermanentWidget(self.status_lbl)
        self.status.addWidget(self.status_lbl)
        self.status_signal.connect(self.update_status)
        self.status_signal.emit('ready', logging.INFO)

    def exit_app(self, checked):
        QApplication.quit()

    def open_file(self):
        mesh_ext = set([".obj", ".stl", '.dae'])
        robot_ext = set([".urdf", ".rvdf"])
        points_ext = set([".txt", ".npy"])
        file = QFileDialog.getOpenFileName(caption="Select a file",
                                           directory=".", filter="Urdf(*.urdf);;Mesh(*.obj *.stl);;PointCloud(*.npy *.txt);;All(*)")[0]
        if not file:
            return
        if Path(file).ext in mesh_ext:
            LoadMesh(file, self.canvas.view, show_edge=True, show_face=True)
            self.status_signal.emit(
                f"{file} has been loaded", logging.INFO)
        elif Path(file).ext in robot_ext:
            self.robot = Robot()
            self.robot.Show(file, self.canvas)
            self.status_signal.emit(
                f"{self.robot.name} has been loaded", logging.INFO)
        elif Path(file).ext in points_ext:
            try:
                pos = np.loadtxt(file)
                if pos.shape[1] == 3:
                    markers = visuals.Markers()
                    markers.set_data(pos, edge_color="grey",
                                     face_color=(1, 0.1, 1, .5), size=3)
                    self.canvas.view.add(markers)
                else:
                    self.status_signal.emit(
                        f"{file} can not be paresed as Nx3 array!", logging.ERROR)
            except Exception as e:
                self.status_signal.emit(
                    f"{file} can not be paresed as Nx3 array!", logging.ERROR)
        self.sidebar.update_tree()

    def update_status(self, msg, level, duration=10000):
        self.status_lbl.setStyleSheet(
            "color:{}".format(self.LOGGER_COLORS[level]))
        self.status_lbl.setText("{}".format((msg)))


class SideBar(QWidget):
    def __init__(self, view, status_signal, app, parent=None):
        super(SideBar, self).__init__(parent)
        vbox = QVBoxLayout()
        self.view = view
        self.status_signal = status_signal
        self.app = app
        self.btn_scene_graph = QPushButton("SceneGraph")
        self.btn_scene_graph.clicked.connect(self.show_tree)
        vbox.addWidget(self.btn_scene_graph)

        self.btn_random = QPushButton("random")
        self.btn_random.clicked.connect(
            lambda: self.app.robot.RandomState(False) if self.app.robot else None)
        vbox.addWidget(self.btn_random)

        self.tree_desc = QLabel(text=view.scene.describe_tree(), parent=self)
        vbox.addWidget(self.tree_desc)

        vbox.addStretch(1)

        self.btn_clear = QPushButton("clear")
        self.btn_clear.clicked.connect(self.clear)
        vbox.addWidget(self.btn_clear)
        self.setLayout(vbox)

    def show_tree(self):
        if self.tree_desc.isVisible():
            self.tree_desc.setVisible(False)
        else:
            self.update_tree()
            self.tree_desc.setVisible(True)

    def update_tree(self):
        self.tree_desc.setText(self.view.scene.describe_tree())

    def clear(self):
        for node in self.view.scene.children:
            if isinstance(node, scene.visuals.VisualNode):
                node.parent = None
        self.view.add(visuals.XYZAxis())
        self.update_tree()
        self.status_signal.emit("clear all visual nodes", logging.INFO)


class RobotControlPanel(QWidget):
    def __init__(self, parent=None):
        super(RobotControlPanel, self).__init__(parent)
        pass


class Canvas(scene.SceneCanvas):

    def __init__(self):
        scene.SceneCanvas.__init__(self, keys=None)
        self.bgcolor = "white"
        self.unfreeze()
        self.view = self.central_widget.add_view()
        self.view.camera = 'turntable'
        # change center to modify rotation center
        # self.view.camera.center = (1, 0, 0)
        self.freeze()
        scene.visuals.XYZAxis(parent=self.view.scene)


def StartApp(*args):
    app = QApplication(*args)
    splash = QSplashScreen()
    splash.setPixmap(QPixmap('resources/robot.png'))
    splash.showMessage('Welcome to RobotZen',
                       Qt.AlignBottom | Qt.AlignCenter, Qt.blue)
    splash.show()
    time.sleep(2)
    window = MainWindow()
    window.show()
    splash.finish(window)
    sys.exit(app.exec_())
